# -*- coding: utf-8 -*- 
"""
  Copyright (c) 2022, colburn
  All rights reserved.
  email:bjay@qq.com
  EM算法示例
"""#"
import argparse
import sys
import matplotlib.pyplot as plt
import time

import numpy as np
import scipy.stats as stats
import torch
import torch.nn as nn
import emtorch as mEm

def loadData(item):
  st ='outfiled' +str(item) +'.npy'
  if torch.cuda.is_available():
    b =torch.Tensor(np.load(st)).cuda()
  else:
    b =torch.Tensor(np.load(st))
  st ='outfilep' +str(item) +'.npz'
  r = np.load(st)
  prt =(r[r.files[0]], r[r.files[1]],
    r[r.files[2]], r[r.files[3]],
    r[r.files[4]]
  )
  le =prt[0].shape[0]
  Ndim =b.shape[0]
  d0 =np.random.random(le)
  d0 =(d0/np.sum(d0)).reshape(le,1)
  
  le =1
  Ndim =b.shape[0]
  d0 =np.random.random(le)
  d0 =(d0/np.sum(d0)).reshape(le,1)
  
  if torch.cuda.is_available():
    dlist =(d0, np.transpose(b[:,0:le].detach().cpu().numpy()).reshape(le,Ndim,1),
      np.ones([le,Ndim,Ndim]) *np.eye(Ndim).reshape(1,Ndim,Ndim),
      np.ones([le,Ndim,Ndim]) *np.eye(Ndim).reshape(1,Ndim,Ndim),
      np.ones([le,1])
    )
  else:
    dlist =(d0, np.transpose(b[:,0:le].numpy()).reshape(le,Ndim,1),
      np.ones([le,Ndim,Ndim]) *np.eye(Ndim).reshape(1,Ndim,Ndim),
      np.ones([le,Ndim,Ndim]) *np.eye(Ndim).reshape(1,Ndim,Ndim),
      np.ones([le,1])
    )
    
    
  return b, prt, dlist

def testReview():
  #  i =0
  for i in range(38):
  #for i in [0,10,15,19,23,25]:
    b, prt, dlist =loadData(i)
    le =dlist[0].shape[0]
    Ndim =b.shape[0]
    em =mEm.myEm()
    markAdd=[]
    for j in range(le):
      markAdd.append(j +1)

    if torch.cuda.is_available():
      device = torch.device("cuda")
      ngpus_per_node = torch.cuda.device_count()
      #b =torch.Tensor(b).cuda()
      dlist =(torch.Tensor(dlist[0]).cuda(), torch.Tensor(dlist[1]).cuda(),
        torch.Tensor(dlist[2]).cuda(), torch.Tensor(dlist[3]).cuda(),
        torch.Tensor(dlist[4]).cuda(), torch.Tensor(markAdd).cuda())
      u =[dlist[1].squeeze(-1).detach().cpu().numpy()]
    else:
      device = torch.device("cpu")
      dlist =(torch.Tensor(dlist[0]), torch.Tensor(dlist[1]),
        torch.Tensor(dlist[2]), torch.Tensor(dlist[3]),
        torch.Tensor(dlist[4]), torch.Tensor(markAdd))
      u =[dlist[1].squeeze(-1).numpy()]
    
    #em.emNdimNormSingleGD(prt, b)
    #td0 =em.forward(b, prt, mu =u, iterations=1000, sFunc =em.emNdimNormSingleGD)
    print(i)
    td0 =em.estimateCluster(b, dlist, mu =u, iterations=1000)
    
    if torch.cuda.is_available():
      c =b.detach().cpu().numpy()
      td =td0[0][1].detach().cpu().numpy()
    else:
      c =b.numpy()
      td =td0[0][1].numpy()
    
    print(td0[2][-1], 'mu', td0[1])
    print(prt,'prt')
    print(td0[0][0], 'td0', i)

    #u =np.asarray(td0[2])
    
    pt =prt[1]
    mEm.drawResult(le, Ndim, c, prt, td, u)
    
if __name__=='__main__':
  testReview()